"""
Correlation
相关性分析模块
"""
from typing import Dict
from typing import Sequence
from typing import Set

import pandas as pd
import pingouin as pg
from scipy.stats import pearsonr

from .base import Score
from .common import DecomposableScorer
from ..model.case import CaseData
from ..model.graph import Graph
from ..model.graph import Node


def partial_correlation(
    node: Node,
    cause: Node,
    graph: Graph,
    series: Dict[Node, Sequence[float]],
    corr_type: str = "pearson",
) -> float:
    """
    Partial correlation coefficient
    偏相关系数计算函数

    corr_type: For the "method" parameter of pingouin.partial_corr
    corr_type: 用于pingouin.partial_corr的"method"参数
    """
    if node == cause:
        return 1  # 如果节点与原因相同，则完全相关，返回1

    # 获取节点和原因的共同父节点作为混淆因素
    confounders = graph.parents(node) | graph.parents(cause)
    confounders -= {node, cause}  # 移除节点自身和原因节点
    data_frame = pd.DataFrame(series)  # 将序列数据转换为DataFrame

    # 获取数据框中的所有节点
    nodes: Set[Node] = set(data_frame.index)
    # 筛选有效的混淆因素：必须在数据中存在且有多个不同的值
    confounders = {
        confounder
        for confounder in confounders
        if confounder in nodes and len(data_frame[confounder].unique()) > 1
    }

    # 使用pingouin库计算偏相关系数
    return pg.partial_corr(
        data=data_frame, x=node, y=cause, covar=confounders, method=corr_type
    )["r"].values[0]


class CorrelationScorer(DecomposableScorer):
    """
    Score nodes by correlation
    通过相关性对节点进行评分的评分器
    """

    def score_node(
        self,
        graph: Graph,
        series: Dict[Node, Sequence[float]],
        node: Node,
        data: CaseData,
    ) -> Score:
        series_node = series[node]  # 获取当前节点的数据序列
        # 计算节点与SLI指标之间的皮尔逊相关系数和p值
        correlation, p_value = pearsonr(series_node, series[data.sli])

        # 创建评分对象，使用相关系数的绝对值作为主评分
        score = Score(abs(correlation))
        score["pearson"] = correlation  # 存储原始相关系数
        score["p-value"] = p_value  # 存储p值，表示统计显著性
        return score


class PartialCorrelationScorer(DecomposableScorer):
    """
    Score nodes by partial correlation coefficient
    通过偏相关系数对节点进行评分的评分器
    """

    def score_node(
        self,
        graph: Graph,
        series: Dict[Node, Sequence[float]],
        node: Node,
        data: CaseData,
    ) -> Score:
        data_frame = pd.DataFrame(series)  # 将序列数据转换为DataFrame
        # 计算SLI指标与当前节点之间的偏相关系数
        correlation = partial_correlation(data.sli, node, graph, data_frame)
        # 创建评分对象，使用偏相关系数的绝对值作为主评分
        score = Score(abs(correlation))
        score["partial-correlation"] = correlation  # 存储原始偏相关系数
        return score
